import numpy as np
import yaml

from ..utils.low_pass_filter import LowPassFilter
from ..robot_state import RobotState


class PVTController:
    def __init__(self, ts: float, config_path: str) -> None:
        super().__init__()

        self.motorName = ['leg_l1_joint', 'leg_l2_joint', 'leg_l3_joint', 'leg_l4_joint', 'leg_l5_joint',
                          'leg_l6_joint', 'leg_r1_joint', 'leg_r2_joint', 'leg_r3_joint', 'leg_r4_joint',
                          'leg_r5_joint', 'leg_r6_joint', ]

        self.jointNum = len(self.motorName)
        self.tau_out_lpf = [LowPassFilter() for _ in range(self.jointNum)]
        self.motor_vel = [0.0 for _ in range(self.jointNum)]
        self.motor_pos_cur = [0.0 for _ in range(self.jointNum)]
        self.motor_pos_des_old = [0.0 for _ in range(self.jointNum)]
        self.motor_tor_out = [0.0 for _ in range(self.jointNum)]
        self.motor_pos_des = [0.0 for _ in range(self.jointNum)]
        self.motor_vel_des = [0.0 for _ in range(self.jointNum)]
        self.motor_tor_des = [0.0 for _ in range(self.jointNum)]

        self.pvt_Kp = [0.0 for _ in range(self.jointNum)]
        self.pvt_Kd = [0.0 for _ in range(self.jointNum)]
        self.maxTor = [400.0 for _ in range(self.jointNum)]
        self.maxVel = [50.0 for _ in range(self.jointNum)]
        self.maxPos = [3.14 for _ in range(self.jointNum)]
        self.minPos = [-3.14 for _ in range(self.jointNum)]

        with open(config_path) as file:
            config = yaml.safe_load(file)

        for i in range(self.jointNum):
            self.pvt_Kp[i] = config[self.motorName[i]]["kp"]
            self.pvt_Kd[i] = config[self.motorName[i]]["kd"]
            self.maxTor[i] = config[self.motorName[i]]["maxTorque"]
            self.maxVel[i] = config[self.motorName[i]]["maxSpeed"]
            self.maxPos[i] = config[self.motorName[i]]["maxPos"]
            self.minPos[i] = config[self.motorName[i]]["minPos"]
            fc = config[self.motorName[i]]["PVT_LPF_Fc"]
            self.tau_out_lpf[i].setPara(fc, ts)
            self.tau_out_lpf[i].ftOut(0)

    def read_state(self, robot_state: RobotState):
        for i in range(self.jointNum):
            self.motor_pos_cur[i] = robot_state.motors_pos_cur[i]
            self.motor_vel[i] = robot_state.motors_vel_cur[i]
        self.motor_pos_des[:] = robot_state.motors_pos_des[:]
        self.motor_vel_des[:] = robot_state.motors_vel_des[:]
        self.motor_tor_des[:] = robot_state.motors_tor_des[:]

    def write_state(self, robot_state: RobotState):
        robot_state.motors_tor_out[:] = self.motor_tor_out[:]
        robot_state.motors_tor_cur[:] = self.motor_tor_out[:]

    def setJointPD(self, kp: float, kd: float, joint_name: str):
        joint_id = self.motorName.index(joint_name)
        self.pvt_Kp[joint_id] = kp
        self.pvt_Kd[joint_id] = kd

    def calMotorsPVT(self, delta_pos_lim=np.inf):
        for i in range(self.jointNum):
            delta = self.motor_pos_des[i] - self.motor_pos_des_old[i]
            if np.abs(delta) > np.abs(delta_pos_lim):
                delta = delta_pos_lim * np.sign(delta)
            p_des = delta + self.motor_pos_des_old[i]
            tau_des = self.pvt_Kp[i] * (p_des - self.motor_pos_cur[i]) + self.pvt_Kd[i] * (
                    self.motor_vel_des[i] - self.motor_vel[i])
            tau_des = self.tau_out_lpf[i].ftOut(tau_des) + self.motor_tor_des[i]
            if np.abs(tau_des) >= np.abs(self.maxTor[i]):
                tau_des = np.sign(tau_des) * self.maxTor[i]
            self.motor_tor_out[i] = tau_des
            self.motor_pos_des_old[i] = p_des
